// ###
// ###  WARNING
// ###
// ###    This file is temporary.
// ###    It does not match the expected public API for BTL 2.
// ###

#include <btl/SLAM/VisualOdometer.hpp>
#include <btl/Maths/OnlineStatistics.hpp>
#include <btl/Utility/OpenCVSupport.hpp>
#include <btl/Utility/EigenSupport.hpp>

#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/features2d/features2d.hpp>
#include <opencv2/calib3d/calib3d.hpp>

#include <Eigen/Geometry>

#include <iostream>
#include <functional>
#include <algorithm>
#include <iterator>
#include <set>

namespace btl
{
namespace slam
{

typedef btl::SmallImage <btl::PixelGrey, OdoPatch::PATCH_SIZE, OdoPatch::PATCH_SIZE> PatchImage;

// ----- OdoFrame -------------------------------------------------

OdoFrame::OdoFrame():
   idx(0) {}

OdoFrame::OdoFrame(const ImageRegionConstGrey& frame, int idx):
   idx(0)
{
   reset(frame, idx);
}

void OdoFrame::reset(const ImageRegionConstGrey& frame, int idx)
{
   reset(frame);
   this->idx = idx;
}

void OdoFrame::reset(const ImageRegionConstGrey& frame)
{
   levels.set(frame, NUM_IMAGE_LEVELS);
}

// ----- OdoPatch -------------------------------------------------

OdoPatch::OdoPatch():
   pos(Eigen::Vector2i::Zero()), level(0), frameIdx(0), score(0.0f)
{
}

OdoPatch::OdoPatch(
   const ImageRegionConstGrey& frame,
   Eigen::Vector2i pos, int level, int frameIdx, float score):
   patch(frame, pos(0), pos(1)), pos(pos), level(level), frameIdx(frameIdx), score(score)
{}

// ----- OdoTrace -------------------------------------------------

int OdoTrace::NextId = 1;

OdoTrace::OdoTrace():
   id(NextId++),
   bestEstimate(Eigen::Vector3d::Zero()),
   matchCount(0), inlierCount(0), outlierCount(0)
{}

// ----- OdoKeyFrame ----------------------------------------------

OdoKeyFrame::OdoKeyFrame():
   pose(Eigen::Affine3d::Identity()), poseFixed(false), hasPoseEstimate(false) {}

// ----- Functors -------------------------------------------------

struct IsOdoPatchDead :
   public std::unary_function<const OdoTrace&, bool>
{
   explicit IsOdoPatchDead(int max_frame_idx):
      max_frame_idx(max_frame_idx) {}

   bool operator()(const OdoTrace& a) const
   {
      return a.lastMatch.frameIdx < max_frame_idx;
   }

   const int max_frame_idx;
};

struct CompareOdoPatchPtrByScoreAscending :
   public std::binary_function<const OdoPatch*, const OdoPatch*, bool>
{
   bool operator()(const OdoPatch* a, const OdoPatch* b) const
   {
      return (a->score < b->score);
   }
};

struct CompareByCellBase
{
   CompareByCellBase(int cell_size):
      cell_size(cell_size) {}

   Eigen::Vector2i cell(const Eigen::Vector2i& v) const
   {
      return vec<int>(v(0) / cell_size, v(1) / cell_size);
   }

   bool cellless(const Eigen::Vector2i& a, const Eigen::Vector2i& b) const
   {
      const Eigen::Vector2i ca = cell(a);
      const Eigen::Vector2i cb = cell(b);
      return
         (ca(1) != cb(1))
         ? (ca(1) < cb(1))  // compare y
         : (ca(0) < cb(0)); // compare x
   }

   bool less(int levelA, int levelB, const Eigen::Vector2i& posA, const Eigen::Vector2i& posB) const
   {
      return
         (levelA != levelB)
         ? (levelA < levelB)
         : cellless(posA, posB);
   }

   const int cell_size;
};

struct CompareOdoPatchByCell :
   private CompareByCellBase,
   public std::binary_function<const OdoPatch&, const OdoPatch&, bool>
{

   CompareOdoPatchByCell(int cell_size):
      CompareByCellBase(cell_size) {}

   bool operator()(const OdoPatch& a, const OdoPatch& b) const
   {
      return less(a.level, b.level, a.pos, b.pos);
   }
};

struct CompareOdoTraceByCell :
   public std::binary_function<const OdoTrace&, const OdoTrace&, bool>
{

   CompareOdoTraceByCell(int cell_size):
      c(cell_size) {}

   bool operator()(const OdoTrace& a, const OdoTrace& b) const
   {
      return c(a.lastMatch, b.lastMatch);
   }

   CompareOdoPatchByCell c;
};

struct CompareCvKeyPointByCell :
   private CompareByCellBase,
   public std::binary_function<const cv::KeyPoint&, const cv::KeyPoint&, bool>
{

   CompareCvKeyPointByCell(int cell_size):
      CompareByCellBase(cell_size) {}

   bool operator()(const cv::KeyPoint& a, const cv::KeyPoint& b) const
   {
      return cellless(vec<int>(a.pt.x, a.pt.y), vec<int>(b.pt.x, b.pt.y));
   }
};

struct CompareKeyFrameMeasurementById :
   public std::binary_function<const OdoKeyFrame::Measurement&, const OdoKeyFrame::Measurement&, bool>
{
   bool operator()(const OdoKeyFrame::Measurement& a, const OdoKeyFrame::Measurement& b) const
   {
      return (a.first < b.first);
   }
};

// ----- VisualOdometer -------------------------------------------

static const int MIN_ACTIVE_LEVEL = 1;
static const int TARGET_CELL_POPULATION = 3;
static const int CELL_SIZE = OdoPatch::PATCH_SIZE * 2;
static const double MAX_MATCH_DISPARITY = 0.04; // 4% of the image width
static const double MIN_MATCH_NCC = 0.7;
static const int MAX_SKIPPED_FRAMES = 3;

static const int FRAMES_PER_KEYFRAME = 15;
static const int KEYFRAMES_PER_REBUILD = 5;
static const int REBUILDS_PER_FIREWALL = 2;

VisualOdometer::VisualOdometer()
{
   _keyframeCountdown = 0;
   // rebuild has to weight for some keyframes
   _rebuildCountdown = KEYFRAMES_PER_REBUILD * FRAMES_PER_KEYFRAME - 1;

   //detector.reset(
   // new cv::GridAdaptedFeatureDetector(
   //    new cv::FastFeatureDetector(4, false),
   //    500, 8, 12));
   _pDetector.reset(new cv::FastFeatureDetector(15, true));

   poses.push_back(Eigen::Affine3d::Identity());
}

VisualOdometer::~VisualOdometer()
{}

void VisualOdometer::setCameraModel(std::auto_ptr<btl::CameraModel>& model)
{
   _pCameraModel.reset(model.release());
}

// conceptually, processing occurs in three nested loops:
//
// while True:
//    for i in range(A):
//       for j in range(B):
//          for k in range(C):
//             detectFeatures()
//             extendTraces()
//             dropDeadTraces()
//             threePointPose()
//             refinePoseEstimate() # motion-only bundle adjustment
//          addNewTraces()
//          addKeyFrame() # record current best pose estimate, if there is one
//          triangulateLandmarks() # triangulation done between two keyframes
//       fivePointPose() # five-point pose done on three keyframes
//       refineKeyFrame() # full bundle adjustment
//       matchSceneScale() # triangulate and match long-running features to fix scene scale
//    insertFirewall()

void VisualOdometer::process(const ImageRegionConstGrey& im)
{
   assert(_pCameraModel);

   frame.reset(im);
   ++frame.idx;
   assert(im.stride() == im.width());

   // always track features and estimate pose
   trackFeatures();
   threePointPose();
   refinePose();

   if(_keyframeCountdown == 0)
   {
      addNewTraces();
      assert(
         std::adjacent_find(traces.rbegin(), traces.rend(),
                            CompareOdoTraceByCell(CELL_SIZE)) == traces.rend());

      // add a new key-frame and re-triangulate observed landmarks
      addKeyFrame();
      triangulateLandmarks();

      _keyframeCountdown = FRAMES_PER_KEYFRAME;
   }

   if(_rebuildCountdown == 0)
   {
      // use 5-point pose on three keyframes to estimate relative pose
      fivePointPose();

      // use pre-emptive RANSAC with 3-point samples to match scale
      // against the previous best estimates of landmark positions
      matchScale();

      _rebuildCountdown = KEYFRAMES_PER_REBUILD * FRAMES_PER_KEYFRAME;
   }

   --_keyframeCountdown;
   --_rebuildCountdown;
}

void VisualOdometer::trackFeatures()
{
   detectFeatures();

   // features must be sorted correctly for extendTraces and addNewTraces to work
   for(int lvl = MIN_ACTIVE_LEVEL; lvl < OdoFrame::NUM_IMAGE_LEVELS; ++lvl)
   {
      assert(
         std::adjacent_find(features[lvl].rbegin(), features[lvl].rend(),
                            CompareOdoPatchByCell(CELL_SIZE)) == features[lvl].rend());
   }

   extendTraces();
   assert(
      std::adjacent_find(traces.rbegin(), traces.rend(),
                         CompareOdoTraceByCell(CELL_SIZE)) == traces.rend());

   dropDeadTraces();
}

void VisualOdometer::addKeyFrame()
{
   std::auto_ptr<OdoKeyFrame> kf(new OdoKeyFrame);

   const int frameIdx = frame.idx;
   kf->frame = frame;

   // copy in measurements from all the traces that were matched this frame
   kf->measurements.reserve(traces.size());
   for(boost::ptr_vector<OdoTrace>::const_iterator
         it = traces.begin(); it != traces.end(); ++it)
   {

      if(! it->patchTrack.empty())
      {
         const OdoTrace::Measurement& m = it->patchTrack.back();
         if(m.first == frameIdx)
         {
            kf->measurements.push_back(
               std::make_pair(it->id, m.second));
         }
      }
   }

   // measurements must be sorted by ID for efficient selection later
   std::sort(
      kf->measurements.begin(),
      kf->measurements.end(),
      CompareKeyFrameMeasurementById());

   // if this is the first keyframe, then its pose is fixed at the origin
   if(keyframes.empty())
   {
      kf->pose = Eigen::Affine3d::Identity();
      kf->poseFixed = true;
      kf->hasPoseEstimate = true;
   }

   std::cout << "Added key frame with " << kf->measurements.size() << " measurements\n";

   // FIXME: initialise with prior estimate from 3-point pose here (but only if we have one)

   keyframes.push_back(kf);
}

void VisualOdometer::threePointPose()
{
}

void VisualOdometer::fivePointPose()
{
   if(keyframes.size() < 3)
      return;

   // run 5-point pose on the first frame and the most recent frame,
   // disambiguating by using whatever frame is in the middle
   OdoKeyFrame& key0 = keyframes.front();
   OdoKeyFrame& key1 = keyframes[keyframes.size() / 2];
   OdoKeyFrame& key2 = keyframes.back();

   typedef Eigen::Matrix<float, 2, 3> Matrix2x3f;
   typedef std::pair<int, Matrix2x3f> Measurement;
   std::vector<Measurement> measurements;

   std::vector<OdoKeyFrame::Measurement>::const_iterator
   it0 = key0.measurements.begin();
   std::vector<OdoKeyFrame::Measurement>::const_iterator
   it1 = key1.measurements.begin();
   std::vector<OdoKeyFrame::Measurement>::const_iterator
   it2 = key2.measurements.begin();

   // select measurements that are present in all three key-frames
   while((it0 != key0.measurements.end()) &&
         (it1 != key1.measurements.end()) &&
         (it2 != key2.measurements.end()))
   {
      int a = it0->first;
      int b = it1->first;
      int c = it2->first;
      if(a == b && b == c)
      {
         Matrix2x3f m;
         m.col(0) = it0->second; ++it0;
         m.col(1) = it1->second; ++it1;
         m.col(2) = it2->second; ++it2;
         measurements.push_back(Measurement(a, m));
      }
      else
      {
         int next = std::max(a, std::max(b, c));
         if(a < next) ++it0;
         if(b < next) ++it1;
         if(c < next) ++it2;
      }
   }

   std::cout << "Performing 5-point pose (" << measurements.size() << " measurements to select from)\n";
}

void VisualOdometer::refinePose()
{
}

void VisualOdometer::triangulateLandmarks()
{
}

void VisualOdometer::matchScale()
{
}

void VisualOdometer::detectFeatures()
{
   std::vector<cv::KeyPoint> pts;

   for(int lvl = MIN_ACTIVE_LEVEL; lvl < OdoFrame::NUM_IMAGE_LEVELS; ++lvl)
   {
      features[lvl].clear();
      pts.clear();
      const ImageRegionConstGrey& im = frame.levels[lvl];
      cv::Mat cvim = cvUnsafeUnconstImage(im);
      _pDetector->detect(cvim, pts);

      // sort detected key-points by cell
      std::sort(pts.begin(), pts.end(),
                CompareCvKeyPointByCell(CELL_SIZE));

      {
         std::vector<cv::KeyPoint>::reverse_iterator
         it = std::adjacent_find(
                 pts.rbegin(), pts.rend(),
                 CompareCvKeyPointByCell(CELL_SIZE));

         assert(it == pts.rend());
      }

      // reserve space for all the detected features
      features[lvl].reserve(pts.size());
      for(std::vector<cv::KeyPoint>::const_iterator
            it = pts.begin(); it != pts.end(); ++it)
      {

         Eigen::Vector2i pos = vec<int>(it->pt.x, it->pt.y);

         const int border = (OdoPatch::PATCH_SIZE + 1) / 2;
         if(pos(0) - border < 0 ||
               pos(1) - border < 0 ||
               pos(0) + border >= im.width() ||
               pos(1) + border >= im.height())
         {
            continue;
         }

         features[lvl].push_back(OdoPatch(frame.levels[lvl], pos, lvl, frame.idx, it->response));
      }
   }
}

void VisualOdometer::extendTraces()
{
   boost::ptr_vector<OdoTrace>::iterator
   tbegin, tend = traces.begin();

   for(int lvl = MIN_ACTIVE_LEVEL; lvl < OdoFrame::NUM_IMAGE_LEVELS; ++lvl)
   {
      tbegin = tend;
      while(tend != traces.end() && tend->lastMatch.level == lvl)
         ++tend;

      extendTraces(frame.levels[lvl], lvl, features[lvl].begin(), features[lvl].end(), tbegin, tend);
      traces.sort(tbegin, tend, CompareOdoTraceByCell(CELL_SIZE));
   }
}

void VisualOdometer::extendTraces(
   const ImageRegionConstGrey& im, int level,
   std::vector<OdoPatch>::const_iterator fbegin,
   std::vector<OdoPatch>::const_iterator fend,
   boost::ptr_vector<OdoTrace>::iterator tbegin,
   boost::ptr_vector<OdoTrace>::iterator tend)
{
   typedef boost::ptr_vector<OdoTrace>::iterator trace_iterator;
   typedef std::vector<OdoPatch>::const_iterator feature_iterator;
   typedef std::pair<trace_iterator, feature_iterator> match;
   std::vector<match> matches;

   const double maxDisparity = MAX_MATCH_DISPARITY * im.width();

   btl::married_matches(
      tbegin, tend,
      std::back_inserter(matches),
      btl::fixed_radius_searcher(fbegin, fend, maxDisparity),
      &btl::features::ncc<OdoPatch::PATCH_SIZE, OdoPatch::PATCH_SIZE>, MIN_MATCH_NCC);

   for(std::vector<match>::iterator it = matches.begin(); it != matches.end(); ++it)
   {
      OdoTrace& trace = *(it->first);
      const OdoPatch& patch = *(it->second);

      trace.lastMatch = patch;

      Eigen::Vector2f posf;
      posf(0) = static_cast<float>(patch.pos(0) << patch.level);
      posf(1) = static_cast<float>(patch.pos(1) << patch.level);
      trace.patchTrack.push_back(std::make_pair(frame.idx, posf));
      ++trace.matchCount;
   }
}

void VisualOdometer::dropDeadTraces()
{
   traces.erase_if(IsOdoPatchDead(frame.idx - MAX_SKIPPED_FRAMES));
}

void VisualOdometer::addNewTraces()
{
   // note: cannot use iterators for this,
   // because traces is modified during the loop
   // (but only by appending, so indexes will remain valid)
   size_t tbegin, tend = 0;
   size_t trace_count = traces.size();

   for(int lvl = MIN_ACTIVE_LEVEL; lvl < OdoFrame::NUM_IMAGE_LEVELS; ++lvl)
   {
      tbegin = tend;
      while(tend != trace_count && traces[tend].lastMatch.level == lvl)
         ++tend;
      addNewTraces(frame.levels[lvl], lvl, features[lvl].begin(), features[lvl].end(), tbegin, tend);
   }

   traces.sort(CompareOdoTraceByCell(CELL_SIZE));
}

void VisualOdometer::addNewTraces(
   const ImageRegionConstGrey& im, int level,
   std::vector<OdoPatch>::const_iterator fbegin,
   std::vector<OdoPatch>::const_iterator fend,
   size_t tbegin, size_t tend)
{

   std::vector<const OdoPatch*> cell_candidates;

   size_t tcellbegin, tcellend = tbegin;
   std::vector<OdoPatch>::const_iterator
   fcellbegin, fcellend = fbegin;

   const int w = im.width();
   const int h = im.height();
   for(int i = 0; i < h; i += CELL_SIZE)
   {
      for(int j = 0; j < w; j += CELL_SIZE)
      {
         Eigen::AlignedBox<int, 2> cell(
            vec(j, i), vec(j + CELL_SIZE - 1, i + CELL_SIZE - 1));

         assert(cell.contains(vec(j, i)));
         assert(!cell.contains(vec(j - 1, i)));
         assert(!cell.contains(vec(j, i - 1)));
         assert(!cell.contains(vec(j + CELL_SIZE, i)));
         assert(!cell.contains(vec(j, i + CELL_SIZE)));

         // get the range of traces in this cell
         tcellbegin = tcellend;
         while(tcellend != tend && cell.contains(traces[tcellend].lastMatch.pos))
            ++tcellend;

         // get the range of features in this cell
         fcellbegin = fcellend;
         while(fcellend != fend && cell.contains(fcellend->pos))
            ++fcellend;

         int count = static_cast<int>(tcellend - tcellbegin);
         if(count < TARGET_CELL_POPULATION)
         {

            size_t remaining_feats = std::distance(fbegin, fend);

            cell_candidates.clear();
            cell_candidates.reserve(std::distance(fcellbegin, fcellend));
            for(std::vector<OdoPatch>::const_iterator
                  it = fcellbegin; it != fcellend; ++it)
               cell_candidates.push_back(&(*it));

            std::sort(
               cell_candidates.begin(),
               cell_candidates.end(),
               CompareOdoPatchPtrByScoreAscending());

            while(count < TARGET_CELL_POPULATION && !cell_candidates.empty())
            {
               const OdoPatch* patch = cell_candidates.back();
               cell_candidates.pop_back();

               std::auto_ptr<OdoTrace> trace(new OdoTrace);
               trace->first = *patch;
               trace->lastMatch = *patch;

               Eigen::Vector2f posf;
               posf(0) = (patch->pos(0) + 0.5f) * (1 << patch->level);
               posf(1) = (patch->pos(1) + 0.5f) * (1 << patch->level);
               trace->patchTrack.push_back(std::make_pair(frame.idx, posf));

               traces.push_back(trace);
               ++count;
            }
         }
         else
         {
            while(fbegin != fend && cell.contains(fbegin->pos))
               ++fbegin;
         }
      }
   }
}

} // namespace slam
} // namespace btl
